__author__ = 'huberg'

try:
    import xml.etree.cElementTree as ET
except ImportError:
    import xml.etree.ElementTree as ET


def load_routes(fname):
    try:
        tree = ET.parse(fname)
    except Exception:
        print "Error: cannot parse file: {0}".format(fname)
        exit(-1)
    root = tree.getroot()
    vechiles = root.findall("vehicle")
    trips = []
    attribs = ['arrivalPos', 'arrivalSpeed', 'depart', 'departLane', 'departPos', 'departSpeed', 'id', 'type']
    for node in vechiles:
        trip = {}
        for k in attribs:
            if k in node.attrib:
                trip.setdefault(k, node.attrib[k])
            else:
                trip.setdefault(k, None)
        # trip.setdefault('route', node[0].attrib['edges'])
        trip.setdefault('route', node[0].attrib['edges'].strip().split())
        trips.append(trip)
    return trips

def load_net(filename):
    tree = ET.ElementTree(file=filename)
    edges = {}
    for node in tree.getroot():
        if node.tag == 'edge':
            edges.setdefault(node.attrib['id'], float(node[0].attrib['length']))
    return edges

def load_vtypes(filename):
    vtypes = []
    for event, eml in ET.iterparse(filename):
        if event == 'end':
            if eml.tag == 'vType':
                vtype = eml.attrib
                if 'probability' not in vtype:
                    print "vtype file error, for each vType, you must specify the probability attrib and the sum of " \
                            "probabilities should be equal to 1"
                    exit(1)
                vtypes.append(vtype)
    return vtypes

def load_group(filename, ignore_header=True):
    """
    :param filename: group file file format : id,label
    :param ignore_header:
    :return: white_group, black_group, type list,
    """
    white_group = []
    black_group = []
    with open(filename, "rb") as f:
        for line in f:
            if ignore_header:
                ignore_header = False
            else:
                vid, label = line.strip().split(',')
                if int(label):
                    white_group.append(vid)
                else:
                    black_group.append(vid)
    return white_group, black_group

def save_group(filename, white_group, black_group, header="id,label"):
    with open(filename, "wb") as f:
        f.write("{0}\n".format(header))
        for item in white_group:
            f.write("{0},1\n".format(item))
        for item in black_group:
            f.write("{0},0\n".format(item))

def write_routes_text(fname, trips):
    attribs = ['arrivalPos', 'arrivalSpeed', 'depart', 'departLane', 'departPos', 'departSpeed', 'id', 'type']
    with open(fname, "wb") as f:
        f.write('<?xml version="1.0" encoding="UTF-8"?>\n')
        f.write('<routes xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:noNamespaceSchemaLocation="http://sumo.dlr.de/xsd/routes_file.xsd">\n')
        for trip in trips:
            line = '\t<vehicle'
            for h in attribs:
                if h in trip and trip[h]:
                    line += ' {0}="{1}"'.format(h, trip[h])
            line +=' >\n'
            f.write(line)
            if 'route' in trip:
                if isinstance(trip['route'], list):
                    f.write('\t\t<route edges="{0}" />\n'.format(" ".join(trip['route'])))
                else:
                    f.write('\t\t<route edges="{0}" />\n'.format(trip['route']))
            if 'stop' in trip:
                assert type(trip['stop']) is dict
                # f.write('\t\t<stop lane="{0}" endPos="{1}" duration="{2}"')
                stop_string = '\t\t<stop'
                for key in trip['stop']:
                    stop_string += ' {0}="{1}"'.format(key, trip['stop'][key])
                stop_string += " />\n"
                f.write(stop_string)
            f.write('\t</vehicle>\n')
        f.write('</routes>\n')
    print 'write routes {0} successful\n'.format(fname)

